package com.kujie.ai.config;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.hc.core5.http.HttpHeaders;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.web.reactive.function.client.WebClientCustomizer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Scope;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;

@Slf4j
@Configuration
public class WebClientConfiguration {
    @Bean
    @Scope("prototype")
    @ConditionalOnMissingBean
    WebClientCustomizer webClientCustomizer() {
        return webClientBuilder -> {
            ExchangeFilterFunction requestFilter = ExchangeFilterFunction.ofRequestProcessor(clientRequest -> {
                log.info("Request Method: {}", clientRequest.method());
                log.info("Request URL: {}", clientRequest.url());
                return Mono.just(clientRequest);
            });

            // 响应拦截器
            ExchangeFilterFunction responseFilter = ExchangeFilterFunction.ofResponseProcessor(clientResponse -> {
                log.info("Intercepted Response Status: {}", clientResponse.statusCode());

                final boolean[] thinking = {false};

                // 处理SSE事件流
                Flux<DataBuffer> modifiedBody = clientResponse.bodyToFlux(DataBuffer.class)
                        .map(buf -> {
                            byte[] bytes = new byte[buf.readableByteCount()];
                            buf.read(bytes);
                            DataBufferUtils.release(buf);
                            return new String(bytes, StandardCharsets.UTF_8);
                        })
                        .flatMap(eventString -> {
                            log.debug(eventString);
                            if (eventString.startsWith("data: ")) {
                                String jsonPart = eventString.substring("data: ".length()).trim();
                                if (JSONUtil.isTypeJSON(jsonPart) && !jsonPart.equals("data: [DONE]")) {
                                    JSONObject retJson;
                                    try {
                                        retJson = JSONUtil.parseObj(jsonPart);
                                    } catch (Exception e) {
                                        log.warn("解析失败");
                                        return Mono.just(eventString);
                                    }

                                    // 修改content字段
                                    JSONArray choices = retJson.getJSONArray("choices");
                                    for (int i = 0; i < choices.size(); i++) {
                                        JSONObject choice = choices.getJSONObject(i);
                                        if (choice == null) {
                                            break;
                                        }
                                        JSONObject delta = choice.getJSONObject("delta");
                                        if (delta == null) {
                                            break;
                                        }
                                        String reasoningContent = delta.getStr("reasoning_content");
                                        String content = delta.getStr("content");
                                        if (StrUtil.isNotBlank(reasoningContent)) {
                                            if (!thinking[0]) {
                                                thinking[0] = true;
                                                delta.set("content", "<think>" + reasoningContent);
                                            } else {
                                                delta.set("content", reasoningContent);
                                            }

                                        } else {
                                            if (thinking[0]) {
                                                thinking[0] = false;
                                                delta.set("content", "</think>" + (content == null ? "" : content));
                                            }
                                        }
                                    }
                                    // 重新生成事件字符串
                                    String modifiedJson = retJson.toString();
                                    return Mono.just("data: " + modifiedJson + "\n\n");
                                }
                                return Mono.just(eventString);
                            }
                            return Mono.just(eventString);
                        })
                        .map(str -> {
                            byte[] bytes = str.getBytes(StandardCharsets.UTF_8);
                            return new DefaultDataBufferFactory().wrap(bytes);
                        });

                // 创建新的ClientResponse，移除CONTENT_LENGTH头
                ClientResponse modifiedResponse = ClientResponse.from(clientResponse)
                        .headers(headers -> headers.remove(HttpHeaders.CONTENT_LENGTH))
                        .body(modifiedBody)
                        .build();

                return Mono.just(modifiedResponse);
            });

            // 将拦截器应用到 WebClient.Builder
            webClientBuilder.filter(requestFilter).filter(responseFilter);
        };
    }
}
